import joblib
from sklearn.feature_extraction.text import TfidfVectorizer
import jieba

# 加载保存的模型和TfidfVectorizer
log_clf = joblib.load('LogisticRegression_model.joblib')
vectorizer = joblib.load('../vectorizer.joblib')

def log_predict_text(text):
    # 使用保存的vectorizer转换新文本
    tfidf = vectorizer.transform([text])
    # 使用加载的模型进行预测
    predicted = log_clf.predict(tfidf)
    return predicted[0]


# 从命令行读取用户输入的文本
if __name__ == '__main__':

    with open('../stopwords.txt', encoding='utf-8') as f:
        con = f.readlines()
        stop_words = set()
        for i in con:
            i = i.replace("\n", "")
            stop_words.add(i)

    while True:

        print("\n请输入新闻（输入'exit'退出）:")
        user_input = input()

        if user_input.lower() == 'exit':
            break
        elif user_input == '':  # 检查是否为空字符串
            print("请重新输入")
            continue

        # 输入处理
        user_input = user_input.replace('\r\n', '').strip()
        user_input = user_input.replace(' ', '').strip()

        result = []
        for word in jieba.lcut(user_input):
            if word not in stop_words:
                result.append(word)

        # 将分词结果列表转换为单个字符串
        text = ' '.join(result)

        # print(text)

        predicted_label = log_predict_text(text)

        # 输出预测结果
        print(f"LogisticRegression模型         预测的类别标签：{predicted_label}")
